package org.xiaoyu.rpc.registry.zookeeper;


import cn.hutool.core.collection.ConcurrentHashSet;
import lombok.extern.slf4j.Slf4j;
import org.xiaoyu.rpc.common.url.URL;
import org.xiaoyu.rpc.registry.Registry;

import java.util.Iterator;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;

/**
 * Created By Have
 * 2021/7/14 22:53
 */
@Slf4j
public class ZookeeperRegistry implements Registry {
    private final ZookeeperInstanceWrapper zkClient;
    private final Set<String> registeredSets = new ConcurrentHashSet<>();
    private final Map<String, URL> registeredURLs = new ConcurrentHashMap<>();

    public ZookeeperRegistry(ZookeeperInstanceWrapper zkClient) {
        this.zkClient = zkClient;
        Runtime.getRuntime().addShutdownHook(new Thread(new Runnable() {
            @Override
            public void run() {
                destroy();
            }
        }));
    }

    @Override
    public void register(URL url) {
        if (registeredSets.contains(url.getAllPath()) || !zkClient.lookupStrictly(url).isEmpty()) {
            registeredSets.add(url.getAllPath());
            registeredURLs.put(url.getAllPath(), url);
            log.info("service node with path [{}] already existed", url.getAllPath());
            return;
        }
        zkClient.registerService(url);
        registeredSets.add(url.getAllPath());
        registeredURLs.put(url.getAllPath(), url);
    }

    @Override
    public void unRegister(URL url) {
        if (!registeredSets.contains(url.getAllPath())) {
            log.info("service node with path [{}] not existed", url.getAllPath());
            return;
        }
        zkClient.unRegistered(url);
        registeredSets.remove(url.getAllPath());
        registeredURLs.remove(url.getAllPath());
    }

    @Override
    public void destroy() {
        Set<Map.Entry<String, URL>> entries = registeredURLs.entrySet();
        Iterator<Map.Entry<String, URL>> iterator = entries.iterator();
        while (iterator.hasNext()) {
            Map.Entry<String, URL> next = iterator.next();
            URL url = next.getValue();
            String key = next.getKey();
            zkClient.unRegistered(url);
            iterator.remove();
            registeredSets.remove(key);
        }
    }
}
